package developer.zyk.ability.chat.ability.rag;

import cn.hutool.core.util.StrUtil;
import developer.zyk.ability.knowledge.KnowledgeClosedLoop;
import developer.zyk.ability.knowledge.bo.OperateRecord;
import developer.zyk.ability.knowledge.common.KnowledgeTypeEnum;
import developer.zyk.ability.knowledge.common.KnowledgeUploadStatus;
import developer.zyk.ability.knowledge.dao.entity.ChatKnowledgeRecordDO;
import developer.zyk.ability.knowledge.parsers.fileParser.FileParserHandler;
import developer.zyk.ability.knowledge.parsers.fileParser.FileParserHandlerFactory;
import developer.zyk.ability.knowledge.parsers.textParser.TextParser;
import developer.zyk.ability.knowledge.service.ChatKnowledgeService;
import developer.zyk.global.common.consts.ChatCommonConstants;
import developer.zyk.global.common.consts.FileType;
import developer.zyk.global.upload.FileOssTemplate;
import developer.zyk.global.upload.FileUtil;
import developer.zyk.global.upload.dao.entity.ChatObjectRepositoryDO;
import developer.zyk.global.user.bo.UserInfo;
import developer.zyk.global.util.AssertUtil;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.document.Document;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.stereotype.Component;

import java.io.File;
import java.util.*;
import java.util.function.BiConsumer;
import java.util.function.Supplier;

@Slf4j
@Component("ragHandleTemplate")
public class RagHandleTemplate {
    @Resource
    private FileParserHandlerFactory fileParserHandlerFactory;
    @Resource
    private FileOssTemplate fileUploadTemplate;
    @Resource
    private VectorStore vectorStore;
    @Resource
    private ChatKnowledgeService chatKnowledgeService;
    @Resource
    private TextParser textParser;


    public void handleKnowledge(Long recordId, UserInfo userInfo) {
        BiConsumer<KnowledgeUploadStatus, String> consumer = (status, msg) -> {
            KnowledgeClosedLoop.record(recordId,
                    new OperateRecord(
                            UUID.randomUUID().toString(),
                            userInfo,
                            status.getCode(),
                            status.getName(),
                            StrUtil.isBlank(msg) ? null : msg
                    ),
                    userInfo,
                    status
            );
        };
        try {
            ChatKnowledgeRecordDO knowledgeRecord = chatKnowledgeService.getById(recordId);
            Boolean enableReset = KnowledgeUploadStatus.queryByCode(chatKnowledgeService.getById(recordId).getStatus()).getEnableParse();
            AssertUtil.assertTrue(enableReset, "该文件不可以进行解析");
            consumer.accept(KnowledgeUploadStatus.START, null);
            Supplier<List<Document>> supplier = getDocument(knowledgeRecord, consumer);
            doHandler(recordId, userInfo, consumer, supplier);
        } catch (Exception e) {
            consumer.accept(
                    KnowledgeUploadStatus.EMBEDDING_FAILED,
                    e.getMessage()
            );
            throw new RuntimeException(e);
        }
    }

    private Supplier<List<Document>> getDocument(ChatKnowledgeRecordDO knowledgeRecord, BiConsumer<KnowledgeUploadStatus, String> statusConsumer) {
        KnowledgeTypeEnum knowledgeTypeEnum = KnowledgeTypeEnum.queryByType(knowledgeRecord.getKnowledgeType());
        switch (knowledgeTypeEnum) {
            case UNSTRUCTURED:
                return () -> {
                    File download = fileUploadTemplate.download(knowledgeRecord.getFileId());
                    ChatObjectRepositoryDO fileInfo = FileUtil.getFileByFileId(knowledgeRecord.getFileId());
                    try {
                        FileParserHandler handler = fileParserHandlerFactory.getHandler(FileType.queryTypeEnum(fileInfo.getFileType()));
                        return handler.parse(download, statusConsumer);
                    } finally {
                        //step 3 任务结束删除文件
                        if (download.exists()) {
                            download.delete();
                            log.debug("文件成功删除,fileId:{}", knowledgeRecord.getFileId());
                        }
                    }
                };
            case STRUCT:
                return () -> {
                    return textParser.parse(knowledgeRecord.getContent());
                };
            default:
                throw new UnsupportedOperationException("暂不不支持解析的知识库类型" + knowledgeTypeEnum.getType());

        }
    }


    private void doHandler(Long recordId, UserInfo userInfo, BiConsumer<KnowledgeUploadStatus, String> consumer, Supplier<List<Document>> documentSupplier) {
        //step 1 解析文件
        List<Document> documents = documentSupplier.get();
        //step 2 注入元数据头
        documents = documents.
                stream()
                .map(a -> {
                    return rebuildMetaData(a, recordId, userInfo);
                })
                .toList();
        //step 3 计算并存储向量
        vectorStore.accept(documents);
        consumer.accept(KnowledgeUploadStatus.VECTOR_STORED_ACCEPT, "");

    }

    private Document rebuildMetaData(Document document, Long recordId, UserInfo userInfo) {
        Map<String, Object> metaData = Optional.of(document.getMetadata()).orElse(new HashMap<>());
        metaData.put(ChatCommonConstants.DocumentMetaDataField.RECORD_ID, recordId);
        metaData.put(ChatCommonConstants.DocumentMetaDataField.USER_ID, userInfo.getUserId());
        return document.mutate()
                .metadata(metaData)
                .build();

    }

}
